package com.hoe.tsb.mmk;

import com.hoe.tsb.tool.SerialUtils;
import com.xiaoleilu.hutool.crypto.digest.DigestUtil;
import lombok.Data;
import lombok.Getter;
import lombok.Setter;
import org.springframework.lang.NonNull;

import java.io.Serializable;
import java.util.*;

/**
 * <p>
 * Multi-Merkle树，添加数据的过程中不允许动态计算叶子节点，只能一次性生成好所有叶子节点，所以不提供add和remove方法
 * </p>
 * 一般树形结构是自顶向下构建，而梅克尔树是自底向上的，而且一旦有新数据加入整棵树都会被重建
 *
 * @author hoe
 * @version 1.0
 * @date 2022/11/28 17:10
 */
public class TraceMmk implements Mmk<TraceData>, Serializable {
    private static final long serialVersionUID = 1L;

    /**
     * 生成的Merkle树
     */
    @Getter
    @Setter
    private Node merkleRoot;
    /**
     * 叶子节点存储数据的Map
     */
    @Getter
    @Setter
    private Map<String, TraceData> leafDataMap;
    /**
     * 所有叶子节点数据的SPV路径
     */
    @Getter
    @Setter
    private Map<String, List<String>> spvMap;

    /**
     * Merkle节点
     */
    @Data
    static class Node implements Serializable {
        private static final long serialVersionUID = 1L;
        Node left;
        Node right;
        Node middle;
        String hash;
    }

    /**
     * 用于FastJSON反序列化用<br/>
     * TODO 后续可以使用专门的序列化/反序列化方法
     */
    public TraceMmk() {

    }

    public TraceMmk(@NonNull List<TraceData> dataList) {
        //一一计算每个item的hash并生成leaf node
        leafDataMap = new HashMap<>(dataList.size());
        spvMap = new HashMap<>(dataList.size());
        Map<String, List<List<String>>> branchSpv = new HashMap<>();
        List<Node> unCalNodes = new ArrayList<>(dataList.size());
        for (TraceData item : dataList) {
            //计算Hash
            String hash = DigestUtil.sha256Hex(SerialUtils.serialize(item));

            //生成叶子节点Node
            Node node = new Node();
            node.hash = hash;
            unCalNodes.add(node);

            //保存数据
            List<String> spvList = new LinkedList<>();
            List<List<String>> spvTmpList = new LinkedList<>();
            spvTmpList.add(spvList);
            leafDataMap.put(hash, item);
            spvMap.put(hash, spvList);
            branchSpv.put(hash, spvTmpList);
        }

        //根据未计算的Node集合构建Merkle树
        while (unCalNodes.size() > 0) {
            for (int idx = 0, size = unCalNodes.size(), odd = size % 2, root = (size == 1 || size == 2 || size == 3) ? 1 : 0; size > 0; idx += 2) {
                Node left = unCalNodes.get(idx), right, middle = null;
                if (idx == 0 && size == 1) {
                    //只有1条数据时，root是和自身进行hash计算出的结果
                    right = left;
                } else if (odd == 1 && size <= 3) {
                    middle = unCalNodes.get(idx + 1);
                    right = unCalNodes.get(idx + 2);
                } else {
                    right = unCalNodes.get(idx + 1);
                }

                Node parent = new Node();
                parent.left = left;
                parent.middle = middle;
                parent.right = right;
                parent.hash = DigestUtil.sha256Hex(left.hash + (middle == null ? "" : middle.hash) + right.hash);
                if (root == 1) {
                    merkleRoot = parent;
                } else {
                    unCalNodes.add(idx, parent);
                }

                //参与计算的元素个数
                int calCnt = 2;
                if (middle != null) {
                    final Node finalMiddle = middle;
                    calCnt = 3;
                    branchSpv.get(left.hash).forEach(l -> l.add(";" + finalMiddle.hash + ";" + right.hash));
                    branchSpv.get(middle.hash).forEach(l -> l.add(left.hash + ";;" + right.hash));
                    branchSpv.get(right.hash).forEach(l -> l.add(left.hash + ";" + finalMiddle.hash + ";"));
                }else if (left == right) {
                    //size为1时出现参与计算的节点只有1个
                    calCnt = 1;
                } else {
                    branchSpv.get(left.hash).forEach(l -> l.add(";" + right.hash));
                    branchSpv.get(right.hash).forEach(l -> l.add(left.hash + ";"));
                }
                rebuildTempMap(branchSpv, left.hash, middle == null ? null : middle.hash, right.hash, parent.hash);

                for (int rCnt = 0, rIdx = root == 1 ? idx : idx + 1; rCnt < calCnt; rCnt++) {
                    unCalNodes.remove(rIdx);
                }
                size -= calCnt;
                idx -= calCnt - 1;
            }
        }
    }

    @Override
    public String getRoot() {
        return merkleRoot.hash;
    }

    @Override
    public TraceData getData(String leafHash) {
        return leafDataMap.get(leafHash);
    }

    @Override
    public List<String> getSpvPath(TraceData traceData) {
        if (traceData == null) {
            return null;
        }

        return spvMap.get(DigestUtil.sha256Hex(SerialUtils.serialize(traceData)));
    }

    @Override
    public boolean contains(String root, TraceData traceData) {
        return getRoot().equals(root) && simpleContains(traceData);
    }

    @Override
    public boolean simpleContains(TraceData traceData) {
        return leafDataMap.get(DigestUtil.sha256Hex(SerialUtils.serialize(traceData))) != null;
    }

    @Override
    public Map<String, TraceData> getLeafDataMap() {
        return leafDataMap;
    }

    @Override
    public boolean equals(Object o) {
        if (o instanceof Mmk<?>) {
            return ((Mmk<?>) o).getRoot().equals(getRoot());
        }
        return false;
    }

    /**
     * 重构临时SPV Map
     * @param map 临时map
     * @param left 当前计算的左节点
     * @param middle 当前计算的中节点
     * @param right 当前计算的右节点
     * @param parent 当前计算的父节点
     */
    private void rebuildTempMap(Map<String, List<List<String>>> map, String left, String middle, String right, String parent) {
        List<List<String>> leftSpvList = map.get(left);
        if (middle != null) {
            leftSpvList.addAll(map.get(middle));
            map.remove(middle);
        }
        leftSpvList.addAll(map.get(right));
        map.remove(left);
        map.remove(right);
        map.put(parent, leftSpvList);
    }

}
